import os, time
import pandas as pd
import numpy as np
# from tqdm import tqdm
import cv2
from pandarallel import pandarallel
pandarallel.initialize()


img_path = '/opt/data/private/project/adc_T9/50UD_05300/selected_train_data/Images'
mask_path = '/opt/data/private/project/adc_T9/50UD_05300/selected_train_data/seg_pred_mask'
# dst_path = '/opt/data/private/project/adc_T9/50UD_05300/selected_train_data/Images_add_mask2'
Wm = 2044
Hm = 2044

df = pd.DataFrame()
codes = []
imgs = []
for code in os.listdir(img_path):
    tmp = os.listdir(os.path.join(img_path, code))
    codes.extend([code] * len(tmp))
    imgs.extend(tmp[:])

# df['root_path'] = [img_path] * len(codes)
df['code'] = codes
df['image'] = imgs

def funcx(code, img_name):
    mask_name = img_name[:-3] + 'png'
    # img = cv2.imread(os.path.join(img_path, code, img_name), 1)
    mask = cv2.imread(os.path.join(mask_path, code, mask_name), 0)
    mask = cv2.resize(mask, (Wm, Hm), interpolation=cv2.INTER_NEAREST)
    x, y = np.where(mask == 1)
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT,(7, 7))
    mask = cv2.dilate(mask,kernel)

    mask[mask == 1.0] = 128
    mask = mask.astype('uint8')
    contours, hierarchy = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    x_list = []
    y_list = []
    radius_list =[]
    # if codes[0] not in args.no_defect_code:
    for cnt in contours:
        (x,y), radius = cv2.minEnclosingCircle(cnt)
        x_list.append(x)
        y_list.append(y)
        radius_list.append(radius)
    if len(radius_list) > 0:
        index = np.argmax(radius_list)
        defect_size_diameter = 2 * np.nanmax(radius_list)
        x, y, w, h = cv2.boundingRect(contours[index])
        return defect_size_diameter, w, h
    else: return 0, 0,0

df['indicators'] = df.parallel_apply(lambda x: funcx(x['code'], x['image']), axis=1) 
df['defect_diameter'] = df['indicators'].apply(lambda x:x[0])
df['defect_w'] = df['indicators'].apply(lambda x:x[1])
df['defect_h']= df['indicators'].apply(lambda x:x[2])
df = df.drop(['indicators'], axis=1)
df.to_csv(r'/opt/data/private/project/adc_T9/50UD_05300/selected_train_data/defect_indicators.csv', index=False)
